import torchvision
from torch.utils .data import DataLoader
from module import *
from time import time
from math import log

train_data = torchvision.datasets.CIFAR10(root=r'./data',train=True,transform=torchvision.transforms.ToTensor(), 
                                          download=True) # 下载训练数据集
test_data = torchvision.datasets.CIFAR10(root=r'./data',train=False,transform=torchvision.transforms.ToTensor(),
                                         download=True) # 下载测试数据集

train_data_size = len(train_data)
test_data_size = len(test_data)
print(f'训练集长度为 {train_data_size}')
print(f'测试集长度为 {test_data_size}')

# 利用dataloader 来加载数据
train_dataloader = DataLoader(train_data, batch_size=64,shuffle=True)
test_dataloader = DataLoader(test_data, batch_size=64,shuffle=True)

# 是否使用gpu
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')

# 创建网络模型
testmodel = TestModule()

# 损失函数
loss_fn = nn.CrossEntropyLoss()

# 优化器
learning_rate = 1e-1
optimizer = torch.optim.SGD(testmodel.parameters(), lr = learning_rate)

# 设置训练的次数
total_train_step = 0
# 设置测试的次数
total_test_step = 0
# 训练的轮数
epoch = 30

# 设置设备
testmodel.to(device=device)
for i in range(1,epoch+1):
    print(f'----------第 {i} 轮训练开始----------')
    testmodel.train()
    start_time = time()
    for imgs,targets in train_dataloader:
        imgs = imgs.to(device)
        targets = targets.to(device)
        output = testmodel(imgs)
        loss = loss_fn(output,targets)
        loss.to(device)
        
        # 优化网络
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()
        total_train_step += 1
        
        if total_train_step % 100 == 0:
            print(f'训练了 {total_train_step} 轮, Loss {loss}')
            end_time = time()
            print('100 轮 用时 {}'.format(end_time-start_time))
            start_time = end_time

    # 测试
    print(f'----------测步骤开始----------')
    total_accuracy = 0
    testmodel.eval()
    with torch.no_grad():
        for imgs, targets in test_dataloader:
            imgs = imgs.to(device)
            targets = targets.to(device)
            outputs = testmodel(imgs)
            preds = outputs.argmax(1)
            accuracy = (preds == targets).sum()
            total_accuracy = total_accuracy + accuracy
            
        print('----------测试正确率 {}----------'.format(total_accuracy/test_data_size))
    
    # 调整学习率
    lr = log(0.9+(total_accuracy/test_data_size/10))/-2*(epoch-i)/epoch
    for group in optimizer.param_groups:
        group['lr'] = lr
    print(f'----------调整学习率为 {lr} ----------')

torch.save(testmodel,r'./weight/0923157.pth')
        
            
